package org.wxd.springcloud.mdc.config.aspect;


import com.google.common.collect.ImmutableMap;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.slf4j.MDC;

import java.util.Map;
import java.util.UUID;

@Aspect
public class ThreadPoolTaskProxy {


    @Pointcut("execution(* org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor.execute(..))")
    private void threadExecuteMethod(){

    }


    @Around("threadExecuteMethod()")
    public void execute(ProceedingJoinPoint pj){
        try {
            Map<String, String> context = MDC.getCopyOfContextMap();

            pj.proceed(new Object[]{wrapRun((Runnable) pj.getArgs()[0],context)});

        } catch (Throwable throwable) {
            throwable.printStackTrace();
        }

    }


    private Runnable wrapRun(Runnable run,Map<String,String> context){
        return () -> {
            String pTraceId = context.get("traceId");
            pTraceId += "|" + UUID.randomUUID().toString();
            MDC.setContextMap(ImmutableMap.of("traceId",pTraceId));
            try {
                run.run();
            }catch (Exception e){

            }finally {
                MDC.clear();
            }
        };



    }

}
